import torch.nn

import torch_plus


class Scaler(torch.nn.Module):
    def __init__(self,
                 input_min: float, input_max: float,
                 output_min: float = -1, output_max: float = 1):
        super().__init__()
        self._input_min = torch_plus.as_tensor(input_min)
        self._input_max = torch_plus.as_tensor(input_max)
        self._output_min = torch_plus.as_tensor(output_min)
        self._output_max = torch_plus.as_tensor(output_max)

        self._k = (self._input_max - self._input_min) / (self._output_max - self._output_min)
        self._b = self._output_min - self._input_min * self._k

    def input_min(self) -> torch.Tensor:
        return self._input_min

    def input_max(self) -> torch.Tensor:
        return self._input_max

    def output_min(self) -> torch.Tensor:
        return self._output_min

    def output_max(self) -> torch.Tensor:
        return self._output_max

    def forward(self, x: torch.Tensor):
        return self._k * x + self._b
